package org.origin.centre.support.base.executor;

import org.origin.centre.support.utils.BasicUtil;
import org.origin.centre.support.utils.TraceUtil;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.concurrent.ListenableFuture;

import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * 自定义执行单元，在执行单元中加上流水号
 *
 * @author ferret
 * @version 2024-05-08
 */
@SuppressWarnings("unused")
public class ThreadPoolTaskExecutorTrace extends ThreadPoolTaskExecutor {

    @Override
    public void execute(Runnable task) {
        super.execute(this.newRunnableMarkLog(task, TraceUtil.getRequestNo()));
    }

    @Override
    public void execute(Runnable task, long startTimeout) {
        super.execute(this.newRunnableMarkLog(task, TraceUtil.getRequestNo()), startTimeout);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(this.newRunnableMarkLog(task, TraceUtil.getRequestNo()));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(this.newCallableMarkLog(task, TraceUtil.getRequestNo()));
    }

    @Override
    public ListenableFuture<?> submitListenable(Runnable task) {
        return super.submitListenable((this.newRunnableMarkLog(task, TraceUtil.getRequestNo())));
    }

    @Override
    public <T> ListenableFuture<T> submitListenable(Callable<T> task) {
        return super.submitListenable(this.newCallableMarkLog(task, TraceUtil.getRequestNo()));
    }

    /**
     * 创建有流水号的 Runnable
     */
    private Runnable newRunnableMarkLog(Runnable runnable, String requestNo) {
        if (BasicUtil.isNotBlank(requestNo)) {
            return () -> {
                try {
                    TraceUtil.addTrace(requestNo);
                    runnable.run();
                } finally {
                    TraceUtil.delTrace();
                }
            };
        }
        return runnable;
    }

    /**
     * 创建有流水号的 Callable
     */
    private <T> Callable<T> newCallableMarkLog(Callable<T> callable, String requestNo) {
        if (BasicUtil.isNotBlank(requestNo)) {
            return () -> {
                try {
                    TraceUtil.addTrace(requestNo);
                    return callable.call();
                } finally {
                    TraceUtil.delTrace();
                }
            };
        }
        return callable;
    }
}
